package com.zhinengshidai.mcp.servers.db;

import com.alibaba.druid.pool.DruidDataSourceFactory;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.zhinengshidai.mcp.servers.common.McpServerManager;
import com.zhinengshidai.mcp.servers.common.annotation.Tool;
import io.modelcontextprotocol.spec.McpSchema;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * Enhanced JDBC Server implementation with improved:
 * 1. Connection pool management
 * 2. Transaction support
 * 3. Batch operation support
 * 4. Better error handling
 * 5. Standardized JSON response format
 *
 * @author chenss
 * @version 3.0.0
 */
public class MysqlMCPServer {
    private static final String SCHEMA_SINGLE_SQL = """
            {
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "SQL query or command"
                    },
                    "params": {
                        "type": "array",
                        "description": "SQL parameters"
                    },
                    "transaction": {
                        "type": "boolean",
                        "description": "Whether to use transaction"
                    }
                },
                "required": ["sql"]
            }
            """;

    private static final String SCHEMA_BATCH_SQL = """
            {
                "type": "object",
                "properties": {
                    "queries": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "sql": {
                                    "type": "string",
                                    "description": "SQL query or command"
                                },
                                "params": {
                                    "type": "array",
                                    "description": "SQL parameters"
                                }
                            },
                            "required": ["sql"]
                        }
                    },
                    "transaction": {
                        "type": "boolean",
                        "description": "Whether to use transaction"
                    }
                },
                "required": ["queries"]
            }
            """;

    private static DataSource dataSource;
    private static volatile boolean initialized = false;

    /**
     * Execute single SQL query/command
     */
    @Tool(description = "execute single mysql SQL query/command, " +
            "return data contains database、table、dataset ," +
            "or no return to insert or update or create", schema = SCHEMA_SINGLE_SQL)
    public String executeSQL(Map<String, Object> params) {
        JSONObject response = new JSONObject();
        response.put("tool", "executeSQL");

        try {
            checkInitialization();
            String sql = params.get("sql").toString();
            boolean transaction = Boolean.parseBoolean(params.getOrDefault("transaction", false).toString());
            Object paramList = params.get("params");
            if (transaction) {
                response.put("result", executeInTransaction(sql, paramList));
            } else {
                response.put("result", executeSingleSQL(sql, paramList));
            }

            response.put("success", true);
        } catch (Exception e) {
            response.put("success", false);
            response.put("error", JSONObject.toJSONString(e));
            logError("executeSQL", e);
        }

        return response.toString();
    }

    /**
     * Execute batch SQL queries/commands
     */
    @Tool(description = "Execute mysql batch SQL queries/commands, return data contains database、table、dataset", schema = SCHEMA_BATCH_SQL)
    public String executeBatchSQL(Map<String, Object> params) {
        JSONObject response = new JSONObject();
        response.put("tool", "executeBatchSQL");

        try {
            checkInitialization();

            JSONArray queries = JSONArray.parseArray(JSONArray.toJSONString(params.get("queries")));
            boolean transaction = Boolean.parseBoolean(params.getOrDefault("transaction", true).toString());

            response.put("result", executeBatch(queries, transaction));
            response.put("success", true);
        } catch (Exception e) {
            response.put("success", false);
            response.put("error", JSONObject.toJSONString(e));
            logError("executeBatchSQL", e);
        }

        return response.toString();
    }


    /**
     * Execute single SQL query/command
     */
    @Tool(description = "check mysql connection state")
    public String checkConnection(Map<String, Object> params) {
        JSONObject response = new JSONObject();
        response.put("tool", "checkConnection");
        try {
            Connection connection = getConnection();
            response.put("result", !connection.isClosed());
        } catch (SQLException e) {
            response.put("result", JSONObject.toJSONString(e));
        }
        return response.toString();
    }

    private static JSONObject executeSingleSQL(String sql, Object paramList) throws SQLException {
        JSONObject result = new JSONObject();

        try (Connection conn = getConnection();
             PreparedStatement pstmt = conn.prepareStatement(sql)) {

            setParameters(pstmt, paramList);

            boolean isQuery = pstmt.execute();
            if (isQuery) {
                try (ResultSet rs = pstmt.getResultSet()) {
                    result.put("data", processResultSet(rs));
                }
            } else {
                result.put("affectedRows", pstmt.getUpdateCount());
            }
        }

        return result;
    }

    private static JSONObject executeInTransaction(String sql, Object paramList) {
        JSONObject result = new JSONObject();
        Connection conn = null;

        try {
            conn = getConnection();
            conn.setAutoCommit(false);

            try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
                setParameters(pstmt, paramList);
                pstmt.execute();
                conn.commit();

                if (pstmt.getResultSet() != null) {
                    try (ResultSet rs = pstmt.getResultSet()) {
                        result.put("data", processResultSet(rs));
                    }
                } else {
                    result.put("affectedRows", pstmt.getUpdateCount());
                }
            }
        } catch (SQLException e) {
            rollbackTransaction(conn, e);
            throw new RuntimeException("Transaction failed", e);
        } finally {
            closeConnection(conn);
        }

        return result;
    }

    private static JSONArray executeBatch(JSONArray queries, boolean transaction) {
        JSONArray results = new JSONArray();
        Connection conn = null;

        try {
            conn = getConnection();
            if (transaction) {
                conn.setAutoCommit(false);
            }

            for (Object queryObj : queries) {
                JSONObject query = JSONObject.parseObject(queryObj.toString());
                JSONObject result = new JSONObject();
                result.put("sql", query.getString("sql"));

                try (PreparedStatement pstmt = conn.prepareStatement(query.getString("sql"))) {
                    setParameters(pstmt, query.get("params"));

                    boolean isQuery = pstmt.execute();
                    if (isQuery) {
                        try (ResultSet rs = pstmt.getResultSet()) {
                            result.put("data", processResultSet(rs));
                        }
                    } else {
                        result.put("affectedRows", pstmt.getUpdateCount());
                    }

                    result.put("success", true);
                } catch (SQLException e) {
                    result.put("success", false);
                    result.put("error", JSONObject.toJSONString(e));

                    if (transaction) {
                        throw e;
                    }
                }

                results.add(result);
            }

            if (transaction) {
                conn.commit();
            }
        } catch (SQLException e) {
            if (transaction) {
                rollbackTransaction(conn, e);
            }
            throw new RuntimeException("Batch execution failed", e);
        } finally {
            if (transaction) {
                closeConnection(conn);
            }
        }

        return results;
    }

    // Helper methods remain largely the same but with improved error handling

    public static void main(String[] args) throws Exception {
        McpServerManager manager = new McpServerManager();
        manager.initStdioServer("mcp-mysql-server", "1.0.0", MysqlMCPServer.class,
                McpSchema.ServerCapabilities.builder()
                        .tools(true)
                        .logging()
                        .build());


    }

    private static String config_prefix = "mysql.";

    private static synchronized void initDataSource() throws Exception {
        if (initialized) return;
        Properties props = new Properties();
        Map<String, String> getenv = System.getenv();
        getenv.forEach((key, value) -> {
            if (key.startsWith(config_prefix)) {
                props.put(key.replace(config_prefix, ""), value);
            }
        });
        dataSource = DruidDataSourceFactory.createDataSource(props);
        initialized = true;
    }

    private static void checkInitialization() {
        if (!initialized) {
            try {
                initDataSource();
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }

    private static void logError(String methodName, Exception e) {
//        System.err.println("Error in " + methodName + ": " + JSONObject.toJSONString(e));
//        e.printStackTrace();
    }

    private static Connection getConnection() throws SQLException {
        checkInitialization();
        return dataSource.getConnection();
    }

    private static void setParameters(PreparedStatement pstmt, Object paramList) throws SQLException {
        if (paramList != null && paramList instanceof List) {
            List<?> params = (List<?>) paramList;
            for (int i = 0; i < params.size(); i++) {
                pstmt.setObject(i + 1, params.get(i));
            }
        }
    }

    private static JSONArray processResultSet(ResultSet rs) throws SQLException {
        JSONArray result = new JSONArray();
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();

        while (rs.next()) {
            JSONObject row = new JSONObject();
            for (int i = 1; i <= columnCount; i++) {
                row.put(metaData.getColumnName(i), rs.getObject(i));
            }
            result.add(row);
        }
        return result;
    }

    private static void rollbackTransaction(Connection conn, Exception e) {
        if (conn != null) {
            try {
                conn.rollback();
            } catch (SQLException ex) {
                logError("rollbackTransaction", ex);
            }
        }
        logError("transaction", e);
    }

    private static void closeConnection(Connection conn) {
        if (conn != null) {
            try {
                if (!conn.getAutoCommit()) {
                    conn.setAutoCommit(true);
                }
                conn.close();
            } catch (SQLException e) {
                logError("closeConnection", e);
            }
        }
    }
}
